#  -*- coding: utf-8 -*-

from data.data_module import DataModule
import traceback
from trading.account import Account
from util.database import DB_CONN
from pathlib import Path
import pandas as pd
from util.stock_util import judge_code_trading_date,get_trading_dates,get_diff_dates,get_sub_industry,get_choice_block_mom_value,get_choice_block_mom_rank
import copy,time
from pymongo import UpdateOne

"""
真实交易数据处理，存入数据库
"""


class LiveTrading:
    def __init__(self,init_cash,init_date):
        self.cash = init_cash
        self.profit = 0.0
        self.init_date = init_date

        self.dm = DataModule()
        self.file = Path(f"{base_code_path}LiveTrading.xlsx")

    def init_trading_record(self):
        trading_record = {
            'code': '',
            'name':'',
            'buy_time':'',
            'sell_time':'',
            'profit':0.0,
            'holding_days': 0,
            'profit_ratio':0.0,
            'buy_avg_price': 0.0,
            'sell_avg_price': 0.0,
            'buy_volume': 0,
            'sell_volume': 0,
            'holding_volume': 0,
            'cost': 0.0,
            'sell_value': 0.0,
            'logic': '',
            'strategy': '',
            'patterns': '',
            'stop_profit_pos': '',
            'stop_profit_price': 0.0,
            'expect_profit_ratio': '',
            'trading_list': list(),
            'analyze': ''
        }
        return trading_record

    def init_date_holding(self):
        date_holding = {
            'date': '',
            'total_value':0.0,
            'date_profit':0.0,
            'date_profit_ratio':0.0,
            'holding_value':0.0,
            'total_profit':self.profit,
            'total_profit_ratio':0.0,
            'holding_profit':0.0,
            'cash':0.0,
            'total_cost':0.0,
            'position': 0.0,
            'code_holding': list()
        }
        return date_holding

    def init_code_holding(self):
        code_holding = {
            'code': '',
            'name':'',
            'code_position':0.0,
            'buy_time':'',
            'cur_time':'',
            'holding_days':0,
            'profit':0.0,
            'profit_ratio':0.0,
            'buy_avg_price': 0.0,
            'sell_avg_price': 0.0,
            'cur_price': 0.0,
            'buy_volume': 0,
            'sell_volume': 0,
            'holding_volume': 0,
            'cost': 0.0,
            'cur_value': 0.0,
            'logic': '',
            'strategy': '',
            'patterns': '',
            'stop_profit_pos': '',
            'stop_profit_price': 0.0,
            'stop_profit_ratio': 0.0,
            'expect_profit_ratio': '',
            'analyze': ''
        }
        return code_holding

    def deal_stock(self,item,record):
        keys = ['code','name','logic','strategy','patterns','stop_profit_pos','stop_profit_price','expect_profit_ratio','analyze']
        for key in keys:
            if item[key]:
                record[key] = item[key]

        record['trading_list'].append(item)

        if record['buy_time'] == '':
            record['buy_time'] = item['date']
        if item['volume'] > 0:
            record['buy_avg_price'] = round((record['buy_avg_price'] * record['buy_volume'] + item['price'] * item['volume'])/(record['buy_volume'] + item['volume']),2)
            record['buy_volume'] = record['buy_volume'] + item['volume']
            record['cost'] = round(record['buy_avg_price'] * record['buy_volume'],2)

        if item['volume'] < 0:
            record['sell_avg_price'] = round((record['sell_avg_price'] * record['sell_volume'] + item['price'] * abs(item['volume']))/(record['sell_volume'] + abs(item['volume'])),2)
            record['sell_volume'] = record['sell_volume'] + abs(item['volume'])
            record['sell_value'] = round(record['sell_avg_price'] * record['sell_volume'],2)

        record['holding_volume'] = record['buy_volume']  - record['sell_volume']


        if record['holding_volume'] == 0:
            record['sell_time'] = item['date']
            #印花税 千分之一
            record['profit'] = round(record['sell_value']  - record['cost'] - round(record['sell_value']/1000,2),2)
            record['profit_ratio'] = round(100*record['profit']/record['cost'],2)
            record['holding_days'] = get_diff_dates(item['code'],record['buy_time'],record['sell_time'])
        return

    def trans_cash(self,item,record):
        #self.cash = round(self.cash + item['price'],2)
        #已经处理过了，不需要了
        return


    def correct_stock(self,item,record):

        return

    def get_code_date_close(self,dm,price,code,date):
        #print(code,date)
        try:
            date_df = dm.get_k_data(code,autype=None,begin_date=date,end_date=date)
            close = date_df.loc[0]['close']
        except:
            close = price
        return close

    def trading_by_close(self,data_df):
        #print(data_df.head())
        data_df['close'] = data_df.apply(lambda row: self.get_code_date_close(self.dm,row['price'],row['code'],row['date']), axis=1)
        data_df['minus_profit'] = data_df.apply(lambda row:round((row['price'] - row['close'])* abs(row['volume']),2),axis=1)
        data_df.to_csv("trading_by_close.csv")

        return

    def import_trading_data(self):

        data_df = pd.read_excel(self.file, encoding="gb2312",dtype=object)
        data_df.rename(columns={'时间': 'date', '代码': 'code', '名称': 'name', '操作类型': 'optype', '操作价格': 'price',
        '操作数量': 'volume', '操作逻辑': 'logic','策略类型': 'strategy','形态学分析':'patterns', '止损位置': 'stop_profit_pos',
        '止损价格': 'stop_profit_price', '盈利预期': 'expect_profit_ratio','后续分析': 'analyze'}, inplace=True)
        data_df['date'] = data_df['date'].astype(str)
        data_df['code'] = data_df['code'].astype(str)
        data_df['strategy'] = data_df['strategy'].astype(str)

        #print(data_df.head())
        update_requests = list()
        collection = DB_CONN['trading_data']
        # 建立code+date的索引，提高save_data时写入数据的查询速度
        collection.create_index([('date', 1)])
        for index, row in data_df.iterrows():

            update_requests.append(
                UpdateOne(
                    {'index': index},
                    {'$set': dict(row)},
                    upsert=True)
            )
        #print(update_requests)
        # 批量写入，提高访问效率
        if len(update_requests) > 0:
            start_time = time.time()
            update_result = collection.bulk_write(update_requests, ordered=False)
            end_time = time.time()
            print('保存交易数据到数据集：%s，插入：%4d条, 更新：%4d条,耗时：%.3f 秒' %
                  (collection.name, update_result.upserted_count, update_result.modified_count,
                   (end_time - start_time)),
                  flush=True)

        return data_df



    def gen_each_trading_record(self,code,code_data,update_requests):
        trading_record = self.init_trading_record()
        op_dict = {
            '买卖':self.deal_stock,
            '转账':self.trans_cash,
            '校正':self.correct_stock,
        }
        for item in code_data:

            op_dict.get(item['optype'], 'wrong value')(item,trading_record)
            #print(f"traingid:{id(trading_record)},trading_record:{trading_record}")
            if trading_record['holding_volume'] == 0 and trading_record['code'] != '':
                update_requests.append(
                    UpdateOne(
                        {'code': code,'buy_time':trading_record['buy_time']},
                        {'$set': trading_record},
                        upsert=True)
                )
                trading_record = self.init_trading_record()

        if trading_record['code'] != '':
            update_requests.append(
                UpdateOne(
                    {'code': code, 'buy_time': trading_record['buy_time']},
                    {'$set': trading_record},
                    upsert=True)
            )
        #print(f"update_requests:{update_requests}")

        return


    def process_trading_record(self):
        collection = DB_CONN['trading_record']
        collection.create_index([('code', 1), ('buy_time', 1)])
        collection.create_index([('sell_time', 1)])

        update_requests = list()

        data = DB_CONN['trading_data'].aggregate([{'$group': {'_id': '$code'}}])
        for item in data:
            code_cursor = DB_CONN['trading_data'].find(
                {'code': item['_id']},
                sort=[('date', 1)],
                projection={'_id': False},
                batch_size=500)
            self.gen_each_trading_record(item['_id'],code_cursor,update_requests)

        #print(update_requests)
        if len(update_requests) > 0:
            start_time = time.time()
            update_result = collection.bulk_write(update_requests, ordered=False)
            end_time = time.time()
            print('保存每一笔交易纪录到数据集：%s，插入：%4d条, 更新：%4d条,耗时：%.3f 秒' %
                  (collection.name, update_result.upserted_count, update_result.modified_count,
                   (end_time - start_time)),
                  flush=True)

        return

    def gen_each_date_holding(self,date,code_cursor,last_date_holding,update_requests):
        date_holding = self.init_date_holding()
        date_holding['date'] = date
        #print(f"date:{date}")
        #code为0，处理转账记录
        org_data_cursor = DB_CONN['trading_data'].find(
            {'code': '0'},
            sort=[('date', 1)],
            projection={'_id': False},
            batch_size=500)

        trans_cash = 0
        for item in org_data_cursor:
            if date < item['date']:
                break
            trans_cash = trans_cash + item['price']

        total_cost = self.cash + trans_cash
        date_holding['total_cost'] = total_cost

        #价值股持仓市值和趋势股持仓市值
        date_holding['value_holding'] = 0
        date_holding['trend_holding'] = 0
        date_holding['other_holding'] = 0

        for item in code_cursor:
            code_holding = self.init_code_holding()
            #print(item['code'],date)
            # if item['code'] == '510500' or item['code'] == '159915':
            #     continue
            trading_date = judge_code_trading_date(item['code'], is_index=False, date=date)
            df_daily = self.dm.get_k_data(item['code'], index=False, autype=None, begin_date=trading_date, end_date=trading_date)
            close = df_daily.loc[0]['close']

            code_holding['code'] = item['code']
            code_holding['name'] = item['name']
            code_holding['buy_time'] = item['buy_time']
            code_holding['cur_time'] = date
            code_holding['holding_days'] = get_diff_dates(item['code'],item['buy_time'],date)
            code_holding['cur_price'] = close
            code_holding['logic'] = item['logic']
            code_holding['stop_profit_pos'] = item['stop_profit_pos']
            code_holding['stop_profit_price'] = item['stop_profit_price']
            code_holding['expect_profit_ratio'] = item['expect_profit_ratio']
            code_holding['analyze'] = item['analyze']
            code_holding['sell_value'] = 0.0
            code_holding['sell_volume'] = 0.0
            code_holding['patterns'] = item['patterns']
            code_holding['strategy'] = item['strategy']
            code_holding['sub_industry'] = get_sub_industry(item['code'])
            code_holding['block_mom_value'] = get_choice_block_mom_value(code_holding['sub_industry'],date)
            code_holding['block_mom_rank'] = get_choice_block_mom_rank(code_holding['sub_industry'],date)

            for trade in item['trading_list']:
                #print(date,type(date),trade['date'],type(trade['date']))
                if date < trade['date']:
                    break
                if trade['volume'] > 0:
                    code_holding['buy_avg_price'] = round(
                        (code_holding['buy_avg_price'] * code_holding['buy_volume'] + trade['price'] * trade['volume'])
                            / (code_holding['buy_volume'] + trade['volume']), 2)
                    code_holding['buy_volume'] = code_holding['buy_volume'] + trade['volume']
                    code_holding['cost'] = round(code_holding['buy_avg_price'] * code_holding['buy_volume'], 2)

                if trade['volume'] < 0:
                    code_holding['sell_avg_price'] = round(
                        (code_holding['sell_avg_price'] * code_holding['sell_volume'] + trade['price'] * abs(trade['volume'])) / (
                                code_holding['sell_volume'] + abs(trade['volume'])), 2)
                    code_holding['sell_volume'] = code_holding['sell_volume'] + abs(trade['volume'])
                    code_holding['sell_value'] = round(code_holding['sell_avg_price'] * code_holding['sell_volume'], 2)

                code_holding['holding_volume'] = code_holding['buy_volume'] - code_holding['sell_volume']
                code_holding['cur_value'] = round(code_holding['holding_volume']*code_holding['cur_price'],2)
                #印花税 千分之一
                code_holding['profit'] = round(code_holding['sell_value'] + code_holding['cur_value'] - code_holding['cost'] - round(code_holding['sell_value']/1000,2), 2)
                code_holding['profit_ratio'] = round(100 * code_holding['profit'] / code_holding['cost'], 2)
                if code_holding['holding_volume'] == 0:
                    self.profit = self.profit + code_holding['profit']

            if code_holding['buy_avg_price'] != 0 and code_holding['stop_profit_price'] != '':
                code_holding['stop_profit_ratio'] = round((code_holding['buy_avg_price'] - code_holding['stop_profit_price'])/code_holding['buy_avg_price'] *100,2)

            if code_holding['strategy'].strip() == "趋势中继":
                date_holding['trend_holding'] += code_holding['cur_value']
            elif code_holding['strategy'].strip() in ["价值定投","低价蓝筹","网格法"]:
                date_holding['value_holding'] += code_holding['cur_value']
            else:
                date_holding['other_holding'] += code_holding['cur_value']



            date_holding['holding_value'] = date_holding['holding_value'] + code_holding['cur_value']
            date_holding['holding_profit'] = date_holding['holding_profit'] + code_holding['profit']

            date_holding['code_holding'].append(code_holding)


        date_holding['total_profit'] = date_holding['total_profit'] + date_holding['holding_profit']
        date_holding['total_value'] = total_cost + date_holding['total_profit']
        date_holding['cash'] = date_holding['total_value'] - date_holding['holding_value']
        date_holding['position'] = round(100 * date_holding['holding_value'] / date_holding['total_value'], 2)
        date_holding['value_position'] = round(100 * date_holding['value_holding'] / date_holding['total_value'], 2)
        date_holding['trend_position'] = round(100 * date_holding['trend_holding'] / date_holding['total_value'], 2)
        date_holding['other_position'] = round(100 * date_holding['other_holding'] / date_holding['total_value'], 2)

        if last_date_holding['total_value'] == 0:
            last_date_holding['total_value'] = total_cost

        date_holding['date_profit'] = round(date_holding['total_value'] - last_date_holding['total_value'],2)
        date_holding['date_profit_ratio'] = round(100* date_holding['date_profit']/last_date_holding['total_value'],2)

        for item in date_holding['code_holding']:
            item['code_position'] = round(100 * item['cur_value'] / date_holding['total_value'], 2)

        update_requests.append(
            UpdateOne(
                {'date': date},
                {'$set': date_holding},
                upsert=True)
        )
        return date_holding


    def process_date_holding(self):
        collection = DB_CONN['date_holding']
        collection.create_index([('code', 1)])
        update_requests = list()
        date_holding = self.init_date_holding()

        dates = get_trading_dates(begin_date=self.init_date)
        for date in dates:

            #找到当天持仓股，当日时间大于买入时间，且小于卖出时间，或卖出时间为空，说明正在持仓
            code_cursor = DB_CONN['trading_record'].find(
                {'buy_time': {'$lte': date},'$or':[{'sell_time': {'$gte': date}},{'sell_time': ''}]},
                sort=[('code', 1)],
                projection={'_id': False},
                batch_size=500)
            date_holding = self.gen_each_date_holding(date,code_cursor,date_holding,update_requests)

        if len(update_requests) > 0:
            start_time = time.time()
            update_result = collection.bulk_write(update_requests, ordered=False)
            end_time = time.time()
            print('保存每一天持仓情况到数据集：%s，插入：%4d条, 更新：%4d条,耗时：%.3f 秒' %
                  (collection.name, update_result.upserted_count, update_result.modified_count,
                   (end_time - start_time)),
                  flush=True)

        return




if __name__ == '__main__':
    live_trading = LiveTrading(305000,'2017-10-10')
    data_df = live_trading.import_trading_data()
    #live_trading.trading_by_close(data_df)
    live_trading.process_trading_record()
    live_trading.process_date_holding()
